# main.py
import os
import argparse
import torch
import time
import numpy as np
import matplotlib.pyplot as plt
from src.dataset import CriteoDataset
from torch.utils.data import DataLoader
from src.model import CriteoMoEModel
from src.trainer import Trainer
from src.utils import set_seed, plot_learning_curves, analyze_gate_weights

# 解决matplotlib中文显示问题
try:
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
except Exception as e:
    print(f"设置中文字体失败: {e}")


def main():
    parser = argparse.ArgumentParser(description='训练Criteo CTR预测的混合专家模型')

    # 数据参数
    parser.add_argument('--data_dir', type=str, default='npy_data', help='预处理后的.npy数据文件目录')
    parser.add_argument('--batch_size', type=int, default=1024, help='批次大小')
    # 建议从一个较低的值开始尝试，例如2
    parser.add_argument('--num_workers', type=int, default=2, help='【训练集】的数据加载工作线程数')
    parser.add_argument('--max_vocab_size', type=int, default=100000, help='分类特征的最大词汇表大小')

    # 模型参数
    parser.add_argument('--embed_dim', type=int, default=8, help='嵌入维度')
    parser.add_argument('--hidden_size', type=int, default=64, help='隐藏层大小')
    # 这里我们使用可以成功运行的自适应协作模型
    # 如果您想回到更早的“无协作”模型，请修改model.py
    # ... (其他模型参数)

    # 训练参数
    parser.add_argument('--learning_rate', type=float, default=0.0001, help='初始学习率')
    parser.add_argument('--weight_decay', type=float, default=0.001, help='权重衰减')
    parser.add_argument('--epochs', type=int, default=10, help='训练轮数')
    parser.add_argument('--early_stopping', type=int, default=3, help='早停耐心值')
    parser.add_argument('--seed', type=int, default=42, help='随机种子')
    parser.add_argument('--log_dir', type=str, default='logs_full_dataset', help='TensorBoard日志目录')
    parser.add_argument('--save_dir', type=str, default='models_full_dataset', help='模型保存目录')
    parser.add_argument('--no_cuda', action='store_true', help='不使用CUDA')

    args = parser.parse_args()

    set_seed(args.seed)
    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(args.save_dir, exist_ok=True)
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    device = torch.device('cuda' if use_cuda else 'cpu')
    print(f"使用设备: {device}")

    try:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        print("--- 加载训练集 ---")
        train_dataset = CriteoDataset(args.data_dir, data_type='train', max_vocab_size=args.max_vocab_size)

        print("--- 加载验证集 ---")
        val_dataset = CriteoDataset(args.data_dir, data_type='val', max_vocab_size=args.max_vocab_size)

        print("--- 加载测试集 ---")
        test_dataset = CriteoDataset(args.data_dir, data_type='test', max_vocab_size=args.max_vocab_size)

        # 关键修改：训练使用多进程，验证/测试使用单进程
        train_loader = DataLoader(
            train_dataset, batch_size=args.batch_size, shuffle=True,
            num_workers=args.num_workers, pin_memory=True, drop_last=True,
            persistent_workers=True if args.num_workers > 0 else False
        )
        val_loader = DataLoader(
            val_dataset, batch_size=args.batch_size * 2, shuffle=False,
            num_workers=0, pin_memory=True
        )
        test_loader = DataLoader(
            test_dataset, batch_size=args.batch_size * 2, shuffle=False,
            num_workers=0, pin_memory=True
        )

        # 模型部分保持不变，使用您最终的自适应协作模型
        model = CriteoMoEModel(
            int_dims=train_dataset.int_dims, cat_dims=train_dataset.cat_dims,
            embed_dim=args.embed_dim, num_experts=2,
            hidden_size=args.hidden_size, gate_hidden=32
        )

        total_params = sum(p.numel() for p in model.parameters())
        print(f"模型总参数数量: {total_params:,}")

        trainer = Trainer(
            model=model, train_loader=train_loader, val_loader=val_loader, test_loader=test_loader,
            learning_rate=args.learning_rate, weight_decay=args.weight_decay,
            log_dir=args.log_dir, save_dir=args.save_dir, device=device
        )

        model, history = trainer.train(num_epochs=args.epochs, early_stopping_patience=args.early_stopping)
        plot_learning_curves(history, save_path=os.path.join(args.save_dir, 'learning_curves.png'))
        trainer.load_best_model()
        test_metrics = trainer.evaluate(test_loader)

        print("\n最终测试集结果 (基于最佳验证模型):")
        print(
            f"Loss: {test_metrics['loss']:.4f}, AUC: {test_metrics['auc']:.4f}, Accuracy: {test_metrics['accuracy']:.4f}, Log Loss: {test_metrics['log_loss']:.4f}")

    except Exception as e:
        print(f"\n训练过程中发生错误: {e}")
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise
    finally:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    print("训练和评估完成!")


if __name__ == "__main__":
    main()